﻿using System;
using System.Drawing;
using System.Drawing.Imaging;

namespace StableDiffusionTools.Imaging
{
    public class Resampling
    {
        public ResamplerType Filter { get; set; }

        public Resampling(ResamplerType filter = ResamplerType.Box)
        {
            Filter = filter;
        }

        /// <summary>
        /// Resamples input array to a new array using current resampling filter.
        /// </summary>
        /// <param name="input">Input array.</param>
        /// <param name="nWidth">Width of the output array.</param>
        /// <param name="nHeight">Height of the output array.</param>
        /// <returns>Output array.</returns>
        public ushort[][,] Resample(ushort[][,] input, int nWidth, int nHeight)
        {

            if (input == null || input.Length == 0 || nWidth <= 1 || nHeight <= 1)
                return null;

            var filter = Resampler.Create(this.Filter);

            int width = input[0].GetLength(0);
            int height = input[0].GetLength(1);
            int planes = input.Length;

            // create bitmaps
            ushort[][,] work = new ushort[planes][,];
            ushort[][,] output = new ushort[planes][,];
            int c = 0;

            for (c = 0; c < planes; c++)
            {

                work[c] = new ushort[nWidth, height];
                output[c] = new ushort[nWidth, nHeight];
            }

            double xScale = ((double)nWidth / width);
            double yScale = ((double)nHeight / height);

            ContributionSum[] contrib = new ContributionSum[nWidth];

            double wdth = 0, center = 0, weight = 0, intensity = 0;
            int left = 0, right = 0, i = 0, j = 0, k = 0;

            // horizontal downsampling
            if (xScale < 1.0)
            {

                // scales from bigger to smaller width
                wdth = (filter.defaultFilterRadius / xScale);

                for (i = 0; i < nWidth; i++)
                {

                    contrib[i].n = 0;
                    contrib[i].p = new Contribution[(int)Math.Floor(2 * wdth + 1)];
                    contrib[i].wsum = 0;
                    center = ((i + 0.5) / xScale);
                    left = (int)(center - wdth);
                    right = (int)(center + wdth);

                    for (j = left; j <= right; j++)
                    {

                        weight = filter.GetValue((center - j - 0.5) * xScale);

                        if ((weight == 0) || (j < 0) || (j >= width))
                            continue;

                        contrib[i].p[contrib[i].n].pixel = j;
                        contrib[i].p[contrib[i].n].weight = weight;
                        contrib[i].wsum += weight;
                        contrib[i].n++;
                    }
                }
            }
            else
            {

                // horizontal upsampling
                // scales from smaller to bigger width
                for (i = 0; i < nWidth; i++)
                {

                    contrib[i].n = 0;
                    contrib[i].p = new Contribution[(int)Math.Floor(2 * filter.defaultFilterRadius + 1)];
                    contrib[i].wsum = 0;
                    center = ((i + 0.5) / xScale);
                    left = (int)Math.Floor(center - filter.defaultFilterRadius);
                    right = (int)Math.Ceiling(center + filter.defaultFilterRadius);

                    for (j = left; j <= right; j++)
                    {

                        weight = filter.GetValue(center - j - 0.5);

                        if ((weight == 0) || (j < 0) || (j >= width))
                            continue;

                        contrib[i].p[contrib[i].n].pixel = j;
                        contrib[i].p[contrib[i].n].weight = weight;
                        contrib[i].wsum += weight;
                        contrib[i].n++;
                    }

                }
            }

            // filter horizontally from input to work
            for (c = 0; c < planes; c++)
            {

                for (k = 0; k < height; k++)
                {

                    for (i = 0; i < nWidth; i++)
                    {

                        intensity = 0;

                        for (j = 0; j < contrib[i].n; j++)
                        {

                            weight = contrib[i].p[j].weight;

                            if (weight == 0)
                                continue;

                            intensity += (input[c][contrib[i].p[j].pixel, k] * weight);
                        }

                        //work[c][i, k] = (ushort)Math.Min(Math.Max(intensity / contrib[i].wsum, UInt16.MinValue), UInt16.MaxValue);
                        work[c][i, k] = (ushort)Math.Min(Math.Max(intensity / contrib[i].wsum + 0.5, UInt16.MinValue), UInt16.MaxValue);
                    }
                }
            }

            // pre-calculate filter contributions for a column
            contrib = new ContributionSum[nHeight];

            // vertical downsampling
            if (yScale < 1.0)
            {

                // scales from bigger to smaller height
                wdth = (filter.defaultFilterRadius / yScale);

                for (i = 0; i < nHeight; i++)
                {

                    contrib[i].n = 0;
                    contrib[i].p = new Contribution[(int)Math.Floor(2 * wdth + 1)];
                    contrib[i].wsum = 0;
                    center = ((i + 0.5) / yScale);
                    left = (int)(center - wdth);
                    right = (int)(center + wdth);

                    for (j = left; j <= right; j++)
                    {

                        weight = filter.GetValue((center - j - 0.5) * yScale);

                        if ((weight == 0) || (j < 0) || (j >= height))
                            continue;

                        contrib[i].p[contrib[i].n].pixel = j;
                        contrib[i].p[contrib[i].n].weight = weight;
                        contrib[i].wsum += weight;
                        contrib[i].n++;
                    }

                }
            }
            else
            {

                // vertical upsampling
                // scales from smaller to bigger height
                for (i = 0; i < nHeight; i++)
                {

                    contrib[i].n = 0;
                    contrib[i].p = new Contribution[(int)Math.Floor(2 * filter.defaultFilterRadius + 1)];
                    contrib[i].wsum = 0;
                    center = ((i + 0.5) / yScale);
                    left = (int)(center - filter.defaultFilterRadius);
                    right = (int)(center + filter.defaultFilterRadius);

                    for (j = left; j <= right; j++)
                    {

                        weight = filter.GetValue(center - j - 0.5);

                        if ((weight == 0) || (j < 0) || (j >= height))
                            continue;

                        contrib[i].p[contrib[i].n].pixel = j;
                        contrib[i].p[contrib[i].n].weight = weight;
                        contrib[i].wsum += weight;
                        contrib[i].n++;
                    }

                }
            }

            // filter vertically from work to output
            for (c = 0; c < planes; c++)
            {

                for (k = 0; k < nWidth; k++)
                {

                    for (i = 0; i < nHeight; i++)
                    {

                        intensity = 0;

                        for (j = 0; j < contrib[i].n; j++)
                        {

                            weight = contrib[i].p[j].weight;

                            if (weight == 0)
                                continue;

                            intensity += (work[c][k, contrib[i].p[j].pixel] * weight);
                        }

                        //output[c][k, i] = (ushort)Math.Min(Math.Max(intensity, UInt16.MinValue), UInt16.MaxValue);
                        output[c][k, i] = (ushort)Math.Min(Math.Max(intensity / contrib[i].wsum + 0.5, UInt16.MinValue), UInt16.MaxValue);
                    }
                }
            }

            work = null;

            return output;
        }


        public static ushort[][,] ConvertBitmapToArray(Bitmap bmp)
        {

            ushort[][,] array = new ushort[4][,];

            for (int i = 0; i < 4; i++)
                array[i] = new ushort[bmp.Width, bmp.Height];


            BitmapData bd = bmp.LockBits(new Rectangle(0, 0, bmp.Width, bmp.Height), ImageLockMode.ReadOnly, PixelFormat.Format32bppArgb);
            //int nOffset = (bd.Stride - bd.Width * 4);
            //int pOffset = (bd.Stride - bd.Width * 1);
            int width = bd.Width;
            int height = bd.Height;
            byte bitsPerPixel = (byte)Image.GetPixelFormatSize(bd.PixelFormat);
            int stride = bd.Stride;
            //byte* scan0 = (byte*)bd.Scan0.ToPointer();

            unsafe
            {
                fixed (ushort* arr3 = array[3])
                fixed (ushort* arr2 = array[2])
                fixed (ushort* arr1 = array[1])
                fixed (ushort* arr0 = array[0])
                {
                    byte* p = (byte*)bd.Scan0.ToPointer();

                    ushort* p0 = arr0;
                    ushort* p1 = arr1;
                    ushort* p2 = arr2;
                    ushort* p3 = arr3;

                    for (int y = 0; y < height; ++y)
                    {
                        for (int x = 0; x < width; ++x)
                        {
                            byte* data = p + y * stride + x * bitsPerPixel / 8;

                            *p0 = data[0];
                            *p1 = data[1];
                            *p2 = data[2];
                            *p3 = data[3];
                            //double magnitude = 1 / 3d * (data[0] + data[1] + data[2]);
                            p0++;
                            p1++;
                            p2++;
                            p3++;
                        }
                    }




                    //    for (int y = 0; y < bd.Height; y++)
                    //    {

                    //        for (int x = 0; x < bd.Width; x++)
                    //        {
                    //            int planeOffset = (y * width) + x;
                    //           p3[planeOffset] = p[3];
                    //           p2[planeOffset] = p[2];
                    //           p1[planeOffset] = p[1];
                    //           p0[planeOffset] = p[0];

                    //            p += 4;
                    //            //p0++;
                    //            //p1++;
                    //            //p2++;
                    //            //p3++;
                    //        }

                    //        //p += nOffset;
                    //        //p0 += nOffset;
                    //        //p1 += nOffset;
                    //        //p2 += nOffset;
                    //        //p3 += nOffset;

                    //        p += nOffset;
                    //    }
                }
            


                //byte* p = (byte*)bd.Scan0;

                //for (int y = 0; y < bd.Height; y++)
                //{
                //    for (int x = 0; x < bd.Width; x++)
                //    {

                //        array[3][x, y] = (ushort)p[3];
                //        array[2][x, y] = (ushort)p[2];
                //        array[1][x, y] = (ushort)p[1];
                //        array[0][x, y] = (ushort)p[0];

                //        p += 4;
                //    }

                //    p += nOffset;
                //}
            }

            bmp.UnlockBits(bd);

            return array;
        }

        /// <summary>
        /// Converts array to Bitmap. Supports only Format32bppArgb pixel format.
        /// </summary>
        /// <param name="array">Array to convert.</param>
        /// <returns>Output Bitmap.</returns>
        public static Bitmap ConvertArrayToBitmap(ushort[][,] array)
        {

            int width = array[0].GetLength(0);
            int height = array[0].GetLength(1);

            Bitmap bmp = new Bitmap(width, height, PixelFormat.Format32bppArgb);

            BitmapData bd = bmp.LockBits(new Rectangle(0, 0, width, height), ImageLockMode.WriteOnly, PixelFormat.Format32bppArgb);
            //int width = bd.Width;
            //int height = bd.Height;
            byte bitsPerPixel = (byte)Image.GetPixelFormatSize(bd.PixelFormat);
            int stride = bd.Stride;

            //int nOffset = (bd.Stride - bd.Width * 4);

            unsafe
            {
                fixed (ushort* arr3 = array[3])
                fixed (ushort* arr2 = array[2])
                fixed (ushort* arr1 = array[1])
                fixed (ushort* arr0 = array[0])
                {
                    byte* p = (byte*)bd.Scan0;
                    ushort* p0 = arr0;
                    ushort* p1 = arr1;
                    ushort* p2 = arr2;
                    ushort* p3 = arr3;

                    for (int y = 0; y < height; y++)
                    {
                        for (int x = 0; x < width; x++)
                        {
                            byte* data = p + y * stride + x * bitsPerPixel / 8;

                            data[0] = (byte)*p0;
                            data[1] = (byte)*p1;
                            data[2] = (byte)*p2;
                            data[3] = (byte)*p3;
                            p0++;
                            p1++;
                            p2++;
                            p3++;

                            //p += 4;
                        }

                        //p += nOffset;
                    }
                }





                //byte* p = (byte*)bd.Scan0;

                //for (int y = 0; y < height; y++)
                //{
                //    for (int x = 0; x < width; x++)
                //    {

                //        p[3] = (byte)Math.Min(Math.Max(array[3][x, y], Byte.MinValue), Byte.MaxValue);
                //        p[2] = (byte)Math.Min(Math.Max(array[2][x, y], Byte.MinValue), Byte.MaxValue);
                //        p[1] = (byte)Math.Min(Math.Max(array[1][x, y], Byte.MinValue), Byte.MaxValue);
                //        p[0] = (byte)Math.Min(Math.Max(array[0][x, y], Byte.MinValue), Byte.MaxValue);

                //        p += 4;
                //    }

                //    p += nOffset;
                //}
            }

            bmp.UnlockBits(bd);

            return bmp;
        }
    }

    internal struct Contribution
    {
        public int pixel;
        public double weight;
    }

    internal struct ContributionSum
    {
        public int n;
        public Contribution[] p;
        public double wsum;
    }

    public enum ResamplerType
    {

        Box = 0,
        Triangle,
        Hermite,
        Bell,
        CubicBSpline,
        Lanczos3,
        Mitchell,
        Cosine,
        CatmullRom,
        Quadratic,
        QuadraticBSpline,
        CubicConvolution,
        Lanczos8
    }


    public abstract class Resampler
    {

        public static Resampler Create(ResamplerType filter)
        {

            Resampler resamplingFilter = null;

            switch (filter)
            {
                case ResamplerType.Box:
                    resamplingFilter = new BoxFilter();
                    break;
                case ResamplerType.Triangle:
                    resamplingFilter = new TriangleFilter();
                    break;
                case ResamplerType.Hermite:
                    resamplingFilter = new HermiteFilter();
                    break;
                case ResamplerType.Bell:
                    resamplingFilter = new BellFilter();
                    break;
                case ResamplerType.CubicBSpline:
                    resamplingFilter = new CubicBSplineFilter();
                    break;
                case ResamplerType.Lanczos3:
                    resamplingFilter = new Lanczos3Filter();
                    break;
                case ResamplerType.Mitchell:
                    resamplingFilter = new MitchellFilter();
                    break;
                case ResamplerType.Cosine:
                    resamplingFilter = new CosineFilter();
                    break;
                case ResamplerType.CatmullRom:
                    resamplingFilter = new CatmullRomFilter();
                    break;
                case ResamplerType.Quadratic:
                    resamplingFilter = new QuadraticFilter();
                    break;
                case ResamplerType.QuadraticBSpline:
                    resamplingFilter = new QuadraticBSplineFilter();
                    break;
                case ResamplerType.CubicConvolution:
                    resamplingFilter = new CubicConvolutionFilter();
                    break;
                case ResamplerType.Lanczos8:
                    resamplingFilter = new Lanczos8Filter();
                    break;
            }

            return resamplingFilter;
        }

        public double defaultFilterRadius;
        public abstract double GetValue(double x);
    }

    internal class HermiteFilter : Resampler
    {

        public HermiteFilter()
        {

            defaultFilterRadius = 1;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            if (x < 1) return (2 * x - 3) * x * x + 1;
            return 0;
        }
    }

    internal class BoxFilter : Resampler
    {

        public BoxFilter()
        {

            defaultFilterRadius = 0.5;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            if (x <= 0.5) return 1;
            return 0;
        }
    }

    internal class TriangleFilter : Resampler
    {

        public TriangleFilter()
        {

            defaultFilterRadius = 1;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            if (x < 1) return 1 - x;
            return 0;
        }
    }

    internal class BellFilter : Resampler
    {

        public BellFilter()
        {

            defaultFilterRadius = 1.5;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            if (x < 0.5) return 0.75 - x * x;
            if (x < 1.5) return 0.5 * Math.Pow(x - 1.5, 2);
            return 0;
        }
    }

    internal class CubicBSplineFilter : Resampler
    {

        double temp;

        public CubicBSplineFilter()
        {

            defaultFilterRadius = 2;
        }

        public override double GetValue(double x)
        {

            if (x < 0)
                x = -x;

            if (x < 1)
            {

                temp = x * x;
                return 0.5 * temp * x - temp + 2f / 3f;
            }
            if (x < 2)
            {

                x = 2f - x;
                return Math.Pow(x, 3) / 6f;
            }
            return 0;
        }
    }

    internal class Lanczos3Filter : Resampler
    {

        public Lanczos3Filter()
        {

            defaultFilterRadius = 3;
        }

        double SinC(double x)
        {

            if (x != 0)
            {

                x *= Math.PI;
                return Math.Sin(x) / x;
            }
            return 1;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            if (x < 3) return SinC(x) * SinC(x / 3f);
            return 0;
        }
    }

    internal class MitchellFilter : Resampler
    {

        const double C = 1 / 3;
        double temp;

        public MitchellFilter()
        {

            defaultFilterRadius = 2;
        }

        public override double GetValue(double x)
        {

            if (x < 0)
                x = -x;

            temp = x * x;
            if (x < 1)
            {

                x = ((12 - 9 * C - 6 * C) * (x * temp)) + ((-18 + 12 * C + 6 * C) * temp) + (6 - 2 * C);
                return (x / 6);
            }
            if (x < 2)
            {

                x = ((-C - 6 * C) * (x * temp)) + ((6 * C + 30 * C) * temp) + ((-12 * C - 48 * C) * x) + (8 * C + 24 * C);
                return (x / 6);
            }
            return 0;
        }
    }

    internal class CosineFilter : Resampler
    {

        public CosineFilter()
        {

            defaultFilterRadius = 1;
        }

        public override double GetValue(double x)
        {

            if ((x >= -1) && (x <= 1)) return (Math.Cos(x * Math.PI) + 1) / 2f;
            return 0;
        }
    }

    internal class CatmullRomFilter : Resampler
    {

        const double C = 1 / 2;
        double temp;

        public CatmullRomFilter()
        {

            defaultFilterRadius = 2;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            temp = x * x;
            if (x <= 1) return 1.5 * temp * x - 2.5 * temp + 1;
            if (x <= 2) return -0.5 * temp * x + 2.5 * temp - 4 * x + 2;
            return 0;
        }
    }

    internal class QuadraticFilter : Resampler
    {

        public QuadraticFilter()
        {

            defaultFilterRadius = 1.5;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            if (x <= 0.5) return -2 * x * x + 1;
            if (x <= 1.5) return x * x - 2.5 * x + 1.5;
            return 0;
        }
    }

    internal class QuadraticBSplineFilter : Resampler
    {

        public QuadraticBSplineFilter()
        {

            defaultFilterRadius = 1.5;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            if (x <= 0.5) return -x * x + 0.75;
            if (x <= 1.5) return 0.5 * x * x - 1.5 * x + 1.125;
            return 0;
        }
    }

    internal class CubicConvolutionFilter : Resampler
    {

        double temp;

        public CubicConvolutionFilter()
        {

            defaultFilterRadius = 3;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            temp = x * x;
            if (x <= 1) return (4f / 3f) * temp * x - (7f / 3f) * temp + 1;
            if (x <= 2) return -(7f / 12f) * temp * x + 3 * temp - (59f / 12f) * x + 2.5;
            if (x <= 3) return (1f / 12f) * temp * x - (2f / 3f) * temp + 1.75 * x - 1.5;
            return 0;
        }
    }

    internal class Lanczos8Filter : Resampler
    {

        public Lanczos8Filter()
        {

            defaultFilterRadius = 8;
        }

        double SinC(double x)
        {

            if (x != 0)
            {
                x *= Math.PI;
                return Math.Sin(x) / x;
            }
            return 1;
        }

        public override double GetValue(double x)
        {

            if (x < 0) x = -x;
            if (x < 8) return SinC(x) * SinC(x / 8f);
            return 0;
        }
    }
}